#include <ros/ros.h>
#include <iostream>
#include <cmath>
#include <Eigen/Dense>
#include "prometheus_msgs/DetectionInfo.h"
#include "prometheus_msgs/ControlCommand.h"
#include "prometheus_msgs/DroneState.h"
#include "prometheus_gimbal_control/GimbalControl.h" 
#include "prometheus_gimbal_control/GimbalState.h"
#include "printf_utils.h"
#include "mission_utils.h"

using namespace std;
using namespace Eigen;

#define NODE_NAME "gimbal_landing"

bool g_use_pad_height; // 是否使用降落板绝对高度
float g_pad_height;
int g_uav_id;
float max_height; // 起始降落位置
float g_camera_offset[3];
float g_distance_to_pad;
float g_arm_height_to_ground;
float g_arm_distance_to_pad;

// include/mission_utils.h 中定义的视觉检测结构体
Detection_result g_landpad_det;   // 检测结果
prometheus_msgs::ControlCommand g_command_now;

prometheus_msgs::DroneState g_UAVState;
Eigen::Matrix3f R_Body_to_ENU;      // 无人机机体系至惯性系旋转矩阵

prometheus_gimbal_control::GimbalState g_gimbal_state;
prometheus_gimbal_control::GimbalControl g_gimbal_cmd;

// 当前无人机飞行模式
std::string curr_mode;

// 消息打印
std::stringstream ss;

// 状态机
enum EXEC_STATE {
    WAITING,
    TRACKING_STAGE1,
    LOST,
    LANDING
};
EXEC_STATE exec_state = WAITING;

// 相机参数
const float lambda = 382;
const float u_max = 319;
const float u_min = -319;
const float v_max = 239;
const float v_min = -239;
const float ud = 120;
const float vd = 120;
const Matrix<float, 2, 4> sd = (Matrix<float, 2, 4>() << -ud, ud, -ud, ud, vd, vd, -vd, -vd).finished();

// 回调函数
void landpadDetCb(const prometheus_msgs::DetectionInfo::ConstPtr &msg) {
    g_landpad_det.object_name = "landpad";
    g_landpad_det.Detection_info = *msg;
    // 识别算法发布的目标位置位于相机坐标系（从相机往前看，物体在相机右方x为正，下方y为正，前方z为正）
    // 相机安装误差 在mission_utils.h中设置场
    // 相机坐标向下，机体系坐标是东北天的设置，x, y轴交换并反向，z轴反向且偏差为负
    // 相机系 -> 机体系
    g_landpad_det.pos_body_frame[0] = -g_landpad_det.Detection_info.position[1] + g_camera_offset[0];
    g_landpad_det.pos_body_frame[1] = -g_landpad_det.Detection_info.position[0] + g_camera_offset[1];
    g_landpad_det.pos_body_frame[2] = -g_landpad_det.Detection_info.position[2] + g_camera_offset[2];

    // 机体系 -> 机体惯性系 (原点在机体的惯性系) (对无人机姿态进行解耦)
    g_landpad_det.pos_body_enu_frame = R_Body_to_ENU * g_landpad_det.pos_body_frame;

    if (g_use_pad_height) {  // 设置是true

        //若已知降落板高度，则无需使用深度信息。
        g_landpad_det.pos_body_enu_frame[2] = g_pad_height - g_UAVState.position[2];
    }

    // 机体惯性系 -> 惯性系，目标在惯性系的位置 = 飞机位置 + 目标在机体惯性系的位置
    g_landpad_det.pos_enu_frame[0] = g_UAVState.position[0] + g_landpad_det.pos_body_enu_frame[0];
    g_landpad_det.pos_enu_frame[1] = g_UAVState.position[1] + g_landpad_det.pos_body_enu_frame[1];
    g_landpad_det.pos_enu_frame[2] = g_UAVState.position[2] + g_landpad_det.pos_body_enu_frame[2];
    // 此降落方案不考虑偏航角
    g_landpad_det.att_enu_frame[2] = 0.0;

    // 目标获得和丢失数积累
    if (g_landpad_det.Detection_info.detected) {
        g_landpad_det.num_regain++;
        g_landpad_det.num_lost = 0;
    }
    else {
        g_landpad_det.num_regain = 0;
        g_landpad_det.num_lost++;
    }

    // 目标丢失，标志位变成false
    if (g_landpad_det.num_lost > VISION_THRES) {
        g_landpad_det.is_detected = false;
    }

    // 目标得到，标志变成true
    if (g_landpad_det.num_regain > VISION_THRES) {
        g_landpad_det.is_detected = true;
    }
}

void droneStateCb(const prometheus_msgs::DroneState::ConstPtr &msg) {
    g_UAVState = *msg;

    curr_mode = msg->mode;

    R_Body_to_ENU = get_rotation_matrix(g_UAVState.attitude[0], g_UAVState.attitude[1], g_UAVState.attitude[2]);

    Eigen::Matrix3f R_c2B = Eigen::Matrix3f::Zero();
    R_c2B(0,1) = -1;
    R_c2B(1,0) = -1;
    R_c2B(2,2) = -1;
}

void gimbalStateCb(const prometheus_gimbal_control::GimbalState::ConstPtr &msg) {
    g_gimbal_state = *msg;
}

// 视觉伺服控制函数
VectorXf controller(const Matrix<float, 2, 4> &uv, const Vector4f &z) {
    // 控制增益参数
    const float k = 5;

    // 视觉框上下限
    Vector4f Mu_over, Mu_under, Mv_over, Mv_under;
    for (int i = 0; i < 4; i++) {
        Mu_over(i) = u_max - sd(0, i);
        Mv_over(i) = v_max - sd(1, i);
        Mu_under(i) = sd(0, i) - u_min;
        Mv_under(i) = sd(1, i) - v_min;
    }

    // 图像雅各比矩阵
    Matrix<float, 8, 6> L = Matrix<float, 8, 6>::Zero();
    for (int i = 0; i < 4; i++) {
        L.block<2, 6>(2 * i, 0) << -lambda / z(i), 0, uv(0, i) / z(i), uv(0, i) * uv(1, i) / lambda, -(lambda * lambda + uv(0, i) * uv(0, i)) / lambda, uv(1, i),
            0, -lambda / z(i), uv(1, i) / z(i), (lambda * lambda + uv(1, i) * uv(1, i)) / lambda, -uv(0, i) * uv(1, i) / lambda, -uv(0, i);
    }

    // 误差
    Vector4f eu = uv.row(0).transpose() - sd.row(0).transpose();
    Vector4f ev = uv.row(1).transpose() - sd.row(1).transpose();

    // Transformed errors
    Vector4f Ku = (1.0 / Mu_over.array() + 1.0 / Mu_under.array()).matrix();
    Vector4f Kv = (1.0 / Mv_over.array() + 1.0 / Mv_under.array()).matrix();

    Vector4f lognu_u = (1.0 + eu.array() / Mu_under.array()).matrix();
    Vector4f lognu_v = (1.0 + ev.array() / Mv_under.array()).matrix();
    Vector4f logde_u = (1.0 - eu.array() / Mu_over.array()).matrix();
    Vector4f logde_v = (1.0 - ev.array() / Mv_over.array()).matrix();

    VectorXf E(8);
    for (int i = 0; i < 4; i++) {
        E(2 * i) = 1.0 / Ku(i) * log(lognu_u(i) / logde_u(i));
        E(2 * i + 1) = 1.0 / Kv(i) * log(lognu_v(i) / logde_v(i));
    }

    // 控制律 - 使用新的伪逆计算方式
    MatrixXf LT = L.transpose();
    MatrixXf L_pinv = (LT * L).inverse() * LT;
    VectorXf Vc = -k * L_pinv * E;
    return Vc;
}

// 云台角速度解析函数
Vector3f gimbalAngularVelo(const VectorXf &Vc, float p, float q, float r, float theta_x, float theta_y, float theta_z) {
    Vector3f Omega_c = Vc.segment<3>(3);
    Vector3f Omega_b(p, q, r);

    Matrix3f M;
    M << 1, 0, -sin(theta_y),
         0, cos(theta_x), sin(theta_x) * cos(theta_y),
         0, -sin(theta_x), cos(theta_x) * cos(theta_y);

    Matrix3f Rx;
    Rx << 1, 0, 0,
          0, cos(theta_x), -sin(theta_x),
          0, sin(theta_x), cos(theta_x);

    Matrix3f Ry;
    Ry << cos(theta_y), 0, sin(theta_y),
          0, 1, 0,
          -sin(theta_y), 0, cos(theta_y);

    Matrix3f Rz;
    Rz << cos(theta_z), -sin(theta_z), 0,
          sin(theta_z), cos(theta_z), 0,
          0, 0, 1;

    Vector3f w_gimbal = M.inverse() * (Omega_c - Rx.transpose() * Ry.transpose() * Rz.transpose() * Omega_b);

    return w_gimbal;
}

// 无人机速度解析函数
Vector3f UAVVelo(const VectorXf &Vc, float theta_x, float theta_y, float theta_z, float phi, float theta, float psi) {
    Vector3f Tc = Vc.segment<3>(0);

    Matrix3f Rx;
    Rx << 1, 0, 0,
          0, cos(phi), -sin(phi),
          0, sin(phi), cos(phi);

    Matrix3f Ry;
    Ry << cos(theta), 0, sin(theta),
          0, 1, 0,
          -sin(theta), 0, cos(theta);

    Matrix3f Rz;
    Rz << cos(psi), -sin(psi), 0,
          sin(psi), cos(psi), 0,
          0, 0, 1;

    Matrix3f Rwb = Rz * Ry * Rx;

    Matrix3f Rxc;
    Rxc << 1, 0, 0,
           0, cos(theta_x), -sin(theta_x),
           0, sin(theta_x), cos(theta_x);

    Matrix3f Ryc;
    Ryc << cos(theta_y), 0, sin(theta_y),
           0, 1, 0,
           -sin(theta_y), 0, cos(theta_y);

    Matrix3f Rzc;
    Rzc << cos(theta_z), -sin(theta_z), 0,
           sin(theta_z), cos(theta_z), 0,
           0, 0, 1;

    Matrix3f Rbc = Rzc * Ryc * Rxc;

    Vector3f V = Rwb * Rbc * Tc;

    return V;
}

// --------------------------- 主函数中inline读取参数 ------------------------------
inline void readParams(const ros::NodeHandle &nh) {
    nh.param<int>("uav_id", g_uav_id, 1);
    //强制上锁高度
    nh.param<float>("arm_height_to_ground", g_arm_height_to_ground, 0.4);
    //强制上锁距离
    nh.param<float>("arm_distance_to_pad", g_arm_distance_to_pad, 0.3);
    // 是否使用降落板绝对高度
    nh.param<bool>("use_pad_height", g_use_pad_height, false);
    nh.param<float>("pad_height", g_pad_height, 0.01);

    // 目标丢失时，最大到飞行高度，如果高于这个值判定为任务失败
    nh.param<float>("max_height", max_height, 3.0);

    // 相机安装偏移,规定为:相机在机体系(质心原点)的位置
    nh.param<float>("camera_offset_x", g_camera_offset[0], 0.0);
    nh.param<float>("camera_offset_y", g_camera_offset[1], 0.0);
    nh.param<float>("camera_offset_z", g_camera_offset[2], 0.0);
}

// ----------------------------------- 主函数 ---------------------------------
int main(int argc, char **argv) {
    ros::init(argc, argv, NODE_NAME);
    ros::NodeHandle nh("~");

    readParams(nh);

    // 【订阅】降落板检测信息
    static ros::Subscriber landpad_det_sub = nh.subscribe<prometheus_msgs::DetectionInfo>("/prometheus/object_detection/landpad_det", 10, landpadDetCb);
    // 【订阅】无人机状态信息
    static ros::Subscriber drone_state_sub = nh.subscribe<prometheus_msgs::DroneState>("/prometheus/drone_state", 10, droneStateCb);
    // 【订阅】云台状态信息
    static ros::Subscriber gimbal_state_sub = nh.subscribe<prometheus_gimbal_control::GimbalState>("/gimbal/state", 10, gimbalStateCb);

    // 【发布】控制命令
    static ros::Publisher g_command_pub = nh.advertise<prometheus_msgs::ControlCommand>("/prometheus/control_command", 10);
    static ros::Publisher g_gimbal_pub = nh.advertise<prometheus_gimbal_control::GimbalControl>("/gimbal/control", 10);

    ros::Rate rate(30.0);

    // 记录时间
    ros::Time stage1_start_time;
    // ros::Time stage2_start_time;
    ros::Time last_time;

    g_command_now.Command_ID = 0;
    while (ros::ok()) {
        ros::spinOnce();
        
        // 持续更新指令ID，不然无法进入OFFBOARD模式
        g_command_now.Command_ID = g_command_now.Command_ID + 1;
        g_command_now.header.stamp = ros::Time::now();
        g_command_now.source = "gimbal_servo_landing";

        // --------------- 未进入OFFBOARD模式，等待 ---------------
        if (curr_mode != "OFFBOARD") {
            exec_state = EXEC_STATE::WAITING;
            ss << "current mode: " << RED << curr_mode << endl
                << BLUE << "Waiting for enter OFFBOARD state \n";
            // PCOUT(1, TAIL, "Waiting for enter OFFBOARD state");
            // continue;
        }

        switch(exec_state) {
            // --------------------- 1. 初始状态，等待视觉检测结果 ----------------------
            case WAITING: {
                // 检测到landingpad就开启降落TRACKING
                if (g_landpad_det.is_detected && curr_mode == "OFFBOARD") {
                    exec_state = TRACKING_STAGE1;

                    // 记录当前时间，开始程序
                    stage1_start_time = ros::Time::now();
                    last_time = ros::Time::now();

                    PCOUT(0, BLUE, "Time to begin.");
                }
                else {
                    // 没有目标时保持云台静止
                    g_gimbal_cmd.rpyMode = g_gimbal_cmd.manual;
                    g_gimbal_cmd.roll = g_gimbal_cmd.noCtl;
                    g_gimbal_cmd.pitch = g_gimbal_cmd.noCtl;
                    g_gimbal_cmd.yaw = g_gimbal_cmd.noCtl;
                    g_gimbal_cmd.rValue = 0;
                    g_gimbal_cmd.pValue = 0;
                    g_gimbal_cmd.yValue = 0;

                    // 添加无人机控制命令
                    // 默认高度为2米，Modules/uav_control/launch/uav_control_outdoor.yaml
                    g_command_now.Mode = prometheus_msgs::ControlCommand::Hold;
                }

                ss << GREEN << "Waiting for the detection result. \n";
                // PCOUT(1, GREEN, "Waiting for the detection result.");
                break;
            }

            // -------------------------- 2. 外层二维码追踪状态 --------------------------
            case TRACKING_STAGE1: {
                if (!g_landpad_det.is_detected) {
                    exec_state = LOST;
                    ROS_WARN("Lost landing pad!");
                    break;
                }

                // 计算视觉控制量
                Matrix<float, 2, 4> uv;
                uv << g_landpad_det.Detection_info.outter_corner0[0], g_landpad_det.Detection_info.outter_corner1[0], g_landpad_det.Detection_info.outter_corner2[0], g_landpad_det.Detection_info.outter_corner3[0],
                    g_landpad_det.Detection_info.outter_corner0[1], g_landpad_det.Detection_info.outter_corner1[1], g_landpad_det.Detection_info.outter_corner2[1], g_landpad_det.Detection_info.outter_corner3[1];
                Vector4f z = Vector4f::Constant(0.5);

                VectorXf Vc = controller(uv, z);

                // 获取无人机当前姿态
                float phi = g_UAVState.attitude[0];
                float theta = g_UAVState.attitude[1];
                float psi = g_UAVState.attitude[2];

                // 获取云台当前姿态
                float theta_x = g_gimbal_state.imuAngle[0];
                float theta_y = g_gimbal_state.imuAngle[1];
                float theta_z = g_gimbal_state.imuAngle[2];

                // 获取无人机当前角速度
                float p = g_UAVState.attitude_rate[0];
                float q = g_UAVState.attitude_rate[1];
                float r = g_UAVState.attitude_rate[2];

                // 计算云台角速度
                Vector3f w_gimbal = gimbalAngularVelo(Vc, p, q, r, theta_x, theta_y, theta_z);

                // 计算无人机速度
                Vector3f V = UAVVelo(Vc, theta_x, theta_y, theta_z, phi, theta, psi);

                // 发布无人机控制指令
                g_command_now.Mode = prometheus_msgs::ControlCommand::Move;
                g_command_now.Reference_State.Move_mode = prometheus_msgs::PositionReference::XYZ_VEL;
                g_command_now.Reference_State.velocity_ref[0] = V(0);
                g_command_now.Reference_State.velocity_ref[1] = V(1);
                g_command_now.Reference_State.velocity_ref[2] = V(2);

                // 发布云台速度控制指令
                g_gimbal_cmd.rpyMode = g_gimbal_cmd.manual;
                g_gimbal_cmd.roll = g_gimbal_cmd.velocityCtl;
                g_gimbal_cmd.pitch = g_gimbal_cmd.velocityCtl;
                g_gimbal_cmd.yaw = g_gimbal_cmd.velocityCtl;
                g_gimbal_cmd.rValue = w_gimbal(0);
                g_gimbal_cmd.pValue = w_gimbal(1);
                g_gimbal_cmd.yValue = w_gimbal(2);

                break;
            }

            // ------------------------- 4. 目标丢失后的任务状态 -----------------------
            case LOST: {
                static int lost_count = 0;
                lost_count++;

                if (lost_count > 30)  // 1秒后仍未找到目标
                {
                    exec_state = LANDING;
                    ROS_ERROR("Unable to find landing pad, initiating landing...");
                }
                else if (g_landpad_det.is_detected) {
                    // exec_state = TRACKING_STAGE1;
                    lost_count = 0;
                    ROS_INFO("Regained landing pad!");
                }
                break;
            }

            // ---------------------------- 5. 降落状态 --------------------------------
            case LANDING: {
                g_command_now.Mode = prometheus_msgs::ControlCommand::Land;
                break;
            }
        }

        g_command_pub.publish(g_command_now);
        g_gimbal_pub.publish(g_gimbal_cmd);

        PCOUT(0.2, TAIL, ss.str());
        ss.str(" ");

        rate.sleep();
    }

    return 0;
}
